import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class ConvBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm=None):
        super(ConvBlock, self).__init__()
        self.conv = torch.nn.Conv2d(input_size, output_size, kernel_size, stride, padding, bias=bias)

        self.norm = norm #根据输入决定归一化方法，卷积层处理后可能使得数据出现一些偏差，这个时候还得跟一个归一化，使得下一次卷积不受影响
        if self.norm =='batch':
            self.bn = torch.nn.BatchNorm2d(output_size)   #批量归一化
        elif self.norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(output_size) #实例归一化

        self.activation = activation #这里是为这个卷积模块增加激活函数看起来有多种可以选择，估计是为了调试方便
        #这些激活函数也得花时间去学啊
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

#这里的代码看起来是先归一化再使用激活函数，这两者的关系是有说法的，可以先归一化，再使用激活函数，也可以先激活函数再归一化，但是会带来不同的效果
    def forward(self, x):
        if self.norm is not None:
            out = self.bn(self.conv(x))
        else:
            out = self.conv(x)

        if self.activation != 'no':
            return self.act(out)
        else:
            return out

class DeconvBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, kernel_size=4, stride=2, padding=1, bias=True, activation='prelu', norm=None):
        super(DeconvBlock, self).__init__()
        #下面这个是torch里面自带的实现反卷积的函数，这个函数默认支持的是零填充（有兴趣探索一下反置卷积）
        self.deconv = torch.nn.ConvTranspose2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
        #同样的归一化
        self.norm = norm
        if self.norm == 'batch':
            self.bn = torch.nn.BatchNorm2d(output_size)
        elif self.norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(output_size)
        #同样的激活函数
        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        if self.norm is not None:
            out = self.bn(self.deconv(x))
        else:
            out = self.deconv(x)

        if self.activation is not None:
            return self.act(out)
        else:
            return out

#解码器部分
class Decoder_MDCBlock1(torch.nn.Module):
    def __init__(self, num_filter, num_ft, kernel_size=4, stride=2, padding=1, bias=True, activation='prelu', norm=None, mode='iter1'):
        super(Decoder_MDCBlock1, self).__init__()
        self.mode = mode
        self.num_ft = num_ft - 1
        self.down_convs = nn.ModuleList()
        self.up_convs = nn.ModuleList()
        for i in range(self.num_ft):
            self.down_convs.append(
                ConvBlock(num_filter*(2**i), num_filter*(2**(i+1)), kernel_size, stride, padding, bias, activation, norm=None)
            )
            self.up_convs.append(
                DeconvBlock(num_filter*(2**(i+1)), num_filter*(2**i), kernel_size, stride, padding, bias, activation, norm=None)
            )
#解码器和编码器的特征融合的区别在于：解码器的输入是一个最浅层的特征，最后残差也是加到这个浅层特征中，输出的也是这个浅层特征。
    #编码器的残差可以理解，如何理解解码器的残差呢？
    #浅层特征升维后我们希望其能提取到下一层的特征，所以要与下一层对比然后反过来作用于浅层特征
    def forward(self, ft_h, ft_l_list):
        if self.mode == 'iter1' or self.mode == 'conv':
            ft_h_list = []
            for i in range(len(ft_l_list)):
                ft_h_list.append(ft_h)
                ft_h = self.down_convs[self.num_ft- len(ft_l_list) + i](ft_h)

            ft_fusion = ft_h
            for i in range(len(ft_l_list)):
                ft_fusion = self.up_convs[self.num_ft-i-1](ft_fusion - ft_l_list[i]) + ft_h_list[len(ft_l_list)-i-1]

        if self.mode == 'iter2':
            ft_fusion = ft_h
            for i in range(len(ft_l_list)):
                ft = ft_fusion
                for j in range(self.num_ft - i):
                    ft = self.down_convs[j](ft)
                ft = ft - ft_l_list[i]
                for j in range(self.num_ft - i):
                    ft = self.up_convs[self.num_ft - i - j - 1](ft)
                ft_fusion = ft_fusion + ft

        if self.mode == 'iter3':
            ft_fusion = ft_h
            for i in range(len(ft_l_list)):
                ft = ft_fusion
                for j in range(i+1):
                    ft = self.down_convs[j](ft)
                ft = ft - ft_l_list[len(ft_l_list) - i - 1]
                for j in range(i+1):
                    # print(j)
                    ft = self.up_convs[i + 1 - j - 1](ft)
                ft_fusion = ft_fusion + ft

        if self.mode == 'iter4':
            ft_fusion = ft_h
            for i in range(len(ft_l_list)):
                ft = ft_h
                for j in range(self.num_ft - i):
                    ft = self.down_convs[j](ft)
                ft = ft - ft_l_list[i]
                for j in range(self.num_ft - i):
                    ft = self.up_convs[self.num_ft - i - j - 1](ft)
                ft_fusion = ft_fusion + ft

        return ft_fusion
#编码器部分
class Encoder_MDCBlock1(torch.nn.Module):
    def __init__(self, num_filter, num_ft, kernel_size=4, stride=2, padding=1, bias=True, activation='prelu', norm=None, mode='iter1'):
        super(Encoder_MDCBlock1, self).__init__()
        #num_filter：输入的通道数
        self.mode = mode #代表特征融合的模式，这里有四种模式，估计是方便调试
        self.num_ft = num_ft - 1 #代表特征融合的级数，好像就是这个编码器有几层的意思
        self.up_convs = nn.ModuleList()#这个函数相当于为网络的模块构建了列表  ，储存上卷积，就是反卷积层的信息
        self.down_convs = nn.ModuleList() #储存正常卷积层的信息
        for i in range(self.num_ft):
            self.up_convs.append(
                #2**i，2的i次方
                DeconvBlock(num_filter//(2**i), num_filter//(2**(i+1)), kernel_size, stride, padding, bias, activation, norm=None)
            )#反卷积处理后通道数变少
            self.down_convs.append(
                ConvBlock(num_filter//(2**(i+1)), num_filter//(2**i), kernel_size, stride, padding, bias, activation, norm=None)
            )#正常卷积处理后通道数变多
            #这个是可以理解的，通道数变多代表正常卷积提取出的特征变多，反卷积通道数变少，代表反卷积在综合特征

    def forward(self, ft_l, ft_h_list):
        #ft_h_list里面存储的是编码器每层提取的高级特征（不用包含第一层特征）
        #ft_l代表的是更高级的特征
        #这里使用的技术有残差连接也有反投影。
        if self.mode == 'iter1' or self.mode == 'conv':
            #融合模式1：大体就是将一个高级特征用反卷积降维后与其对应差一级的特征比较得到差距，将差距升维后与高级特征相加，然后返回
            #因为浅层细节到抽象语义之间的转换必然存在细节的丢失，这样设计的目的是为了能够尽可能弥补丢失的信息
            ft_l_list = []
            for i in range(len(ft_h_list)):
                ft_l_list.append(ft_l)
                ft_l = self.up_convs[self.num_ft- len(ft_h_list) + i](ft_l)#up_convs本身是反卷积的一个列表这里是按照索引访问了对应的反卷积层然后进行一个升维操作
              #这里索引中的self.num_ft- len(ft_h_list)其实是为了方便动态训练在一些情况下可能不需要融合全部的特征，这里索引设置的其实很巧妙，有兴趣可以探究一下
            ft_fusion = ft_l
            for i in range(len(ft_h_list)):
                ft_fusion = self.down_convs[self.num_ft-i-1](ft_fusion - ft_h_list[i]) + ft_l_list[len(ft_h_list)-i-1]

        if self.mode == 'iter2':
            #融合模式2:
            #融合模式2相比融合模式1更加的复杂和强大
            #这个模式2就是对于每个比最高级特征浅的特征都进行一次操作，比如对于最浅近的特征，先将最高级特征降维到对应的特征尺寸，做差后再逐步升维到最高级特征尺寸然后做和，每个特征都要这样循环一遍
            ft_fusion = ft_l
            for i in range(len(ft_h_list)):
                ft = ft_fusion
                for j in range(self.num_ft - i):
                    ft = self.up_convs[j](ft)
                ft = ft - ft_h_list[i]
                for j in range(self.num_ft - i):
                    # print(j)
                    ft = self.down_convs[self.num_ft - i - j - 1](ft)
                ft_fusion = ft_fusion + ft
        #融合模式1和融合模式2在实现上有差别，最后得到的结果也有很大差别，之后可以思考一下这种差别源自什么
        #AI说融合模式3结合了融合模式1和2，是最巧妙的一个融合方式
        #融合模式3和2很像最大的区别就是2的特征融合永远在最高级特征层次，而3在中间层次完成特征融合

        if self.mode == 'iter3':
            ft_fusion = ft_l
            for i in range(len(ft_h_list)):
                ft = ft_fusion
                for j in range(i+1):
                    ft = self.up_convs[j](ft)
                ft = ft - ft_h_list[len(ft_h_list) - i - 1]
                for j in range(i+1):
                    # print(j)
                    ft = self.down_convs[i + 1 - j - 1](ft)
                ft_fusion = ft_fusion + ft
        #融合模式4是最直接的融合方式
        #融合模式4和融合模式2看起来很像，最主要差别是2的融合是从每次融合的结果上得来的，是叠加的，4是每次都从基本的最高级特征出发是分开来的，各个过程互不影响

        if self.mode == 'iter4':
            ft_fusion = ft_l
            for i in range(len(ft_h_list)):
                ft = ft_l
                for j in range(self.num_ft - i):
                    ft = self.up_convs[j](ft)
                ft = ft - ft_h_list[i]
                for j in range(self.num_ft - i):
                    # print(j)
                    ft = self.down_convs[self.num_ft - i - j - 1](ft)
                ft_fusion = ft_fusion + ft

        return ft_fusion
#看完编码器和解码器的结构，我发现和我想象的有点不太一样啊。编码器一直致力于精简特征，正向的一个驱动
#解码器致力于尽可能通过特征还原图像，是一个逆向的过程，从正反两个方向同时出发，有点意思